import pandas as pd
from io import StringIO
import datetime

class DataHandler:
    """
    数据模块 (Data Module) - V2 (已更新以包含ETF价格)

    核心职责:
    1. 从本地CSV文件加载行情、指标以及新增的ETF数据。
    2. 对数据进行清洗和标准化（如：处理日期、重命名列）。
    3. 提供一个统一的数据访问接口，能够根据指定的日期范围，
       以生成器(generator)的形式逐日向回测核心输送数据。
    """

    def __init__(self, csv_filepath):
        """
        初始化数据处理器。
        
        在实例化时，会立即加载并预处理CSV文件。
        
        :param csv_filepath: CSV文件的路径 (可以是文件路径字符串，或类似文件的对象)。
        """
        self.csv_filepath = csv_filepath
        self.data_df = self._load_and_prepare_data()

    def _load_and_prepare_data(self):
        """
        一个私有方法，用于从CSV文件加载数据并进行预处理。
        
        处理步骤:
        - 使用pandas读取CSV。
        - 将关键的中文列名映射为标准的英文名，方便代码调用。
        - 将'日期'列转换为datetime对象，并将其设置为DataFrame的索引。
        - 确保整个数据表是按日期升序排列的。
        
        :return: 经过处理后的pandas DataFrame。
        """
        try:
            # 从文件路径加载数据
            df = pd.read_csv(self.csv_filepath)
            
            # --- 数据标准化 ---
            # 1. 定义列名映射关系 (中文 -> 英文)，便于程序化访问
            # 【核心修改】在这里增加了 '沪深300ETF' 的映射
            column_mapping = {
                '日期': 'date',
                '开盘': 'open',
                '收盘': 'close',
                '最高': 'high',
                '最低': 'low',
                '成交量': 'volume',
                '成交额': 'turnover',
                '振幅': 'amplitude',
                '涨跌幅': 'pct_change',
                '涨跌额': 'change',
                '换手率': 'turnover_rate',
                '10年期国债收益率': 'bond_yield',
                '滚动市盈率': 'pe_ratio_ttm',
                '股债利差': 'spread',
                '沪深300ETF': 'etf_close'  # <--- 新增的映射
            }
            df.rename(columns=column_mapping, inplace=True)

            # 2. 转换日期列为标准datetime格式，并设为索引
            df['date'] = pd.to_datetime(df['date'])
            df.set_index('date', inplace=True)
            
            # 3. 确保数据是按时间顺序排列的
            df.sort_index(ascending=True, inplace=True)

            print("数据加载和预处理成功（已包含ETF数据）。")
            return df
        except FileNotFoundError:
            print(f"错误：无法找到文件 {self.csv_filepath}")
            return pd.DataFrame()
        except KeyError as e:
            print(f"错误：CSV文件中缺少必要的列: {e}。请检查文件是否包含所有必需的列名，包括 '沪深300ETF'。")
            return pd.DataFrame()
        except Exception as e:
            print(f"加载或处理数据时发生未知错误: {e}")
            return pd.DataFrame()

    def get_data_generator(self, start_date, end_date):
        """
        根据指定的开始和结束日期，创建一个生成器来逐日提供数据。
        这是数据模块提供给回测核心的主要接口。
        
        :param start_date: 回测开始日期 (字符串 'YYYY-MM-DD' 或 datetime 对象)
        :param end_date: 回测结束日期 (字符串 'YYYY-MM-DD' 或 datetime 对象)
        :return: 一个生成器。如果范围内有数据，它会逐一产出(yield)一个命名元组(named tuple)，
                 其中包含了当天的所有数据。否则返回None。
        """
        # --- 数据过滤 ---
        # 使用DataFrame索引进行高效的日期范围切片
        filtered_df = self.data_df.loc[start_date:end_date]

        if filtered_df.empty:
            print(f"警告：在 {start_date} 到 {end_date} 的日期范围内没有找到任何数据。")
            return

        
        # --- 数据接口 ---
        # 使用 itertuples() 创建一个高效的迭代器。
        # 它返回的是命名元组，比 iterrows() 更快，访问也更方便。
        # 回测核心将通过 for 循环来逐日获取数据。
        for row in filtered_df.itertuples():
            yield row
